# -*- coding: utf-8 -*-
# @desc: 根据回测任务，生成回测参数和回测代码。如果有多个子任务，则生成多份。
# @file: code_generator.py
# @author: Lin Jackey Lyu (me@linlyu.com)
# @date: 2018/9/28 17:45
from enum import Enum
import datetime
import os
from core.config import YamlConfig
from .strategy_template import HOLD_AND_BUY_TEMPLATE, ORDERS_TEMPLATE, \
    REBALANCE_TEMPLATE, HOLD_SPEC_DAYS_TEMPLATE


class PolicyClass(Enum):
    hold_and_buy = 1
    orders = 2
    rebalance = 3
    hold_spec_days = 4


class BTCodeGenerator(object):

    @staticmethod
    def generate(config: YamlConfig, selector, policy_type) -> (list, None):
        start_date = config.spec.start_date
        end_date = config.spec.end_date 
        account = config.spec.get('account', 100000)
        duration = config.spec.get('duration', 2)
        hold_cnt = config.spec.get('hold_cnt', 5)
        gain_limit1 = config.spec.get('gain_limit1', None)
        gain_limit2 = config.spec.get('gain_limit2', None)
        loss_limit = config.spec.get('loss_limit', None)
        signal = config.spec.get('signal', False)
        price_limit = config.spec.get("price_limit", True)

        data_bundle_path = config.rqalpha.bundle_data_path

        bt_config = {
            "base": {
                "benchmark": '000300.XSHG',  # 沪深300指数
                'data_bundle_path': os.path.abspath(os.path.expanduser(data_bundle_path)),
                "frequency": "1d",
                "start_date": start_date,
                "end_date": end_date,
                "accounts": {
                    "stock": account  # 初始资金
                }
            },
            "mod": {
                "sys_simulation": {
                    # 是否开启信号模式
                    "signal": signal,
                    "enabled": True,
                    # "slippage": 0.01
                    # price_limit: 在处于涨跌停时，无法买进/卖出，默认开启【在 Signal 模式下，不再禁止买进/卖出，如果开启，则给出警告提示。】
                    "price_limit": price_limit,
                }
            },
            "extra": {
                "log_level": "info",
                "logger": [("user_log", config.spec.logging)],
                "context_vars": {
                    "selector": selector  # 选股逻辑的实现
                }
            }
        }

        assert hold_cnt > 0, 'hold_cnt({}) is <= 0.'.format(hold_cnt)
        assert duration > 0, 'duration({}) is <=0.'.format(duration)
        # assert gain_limit1 > gain_limit2 > 0
        assert loss_limit < 0

        if policy_type == PolicyClass.hold_and_buy:
            code = HOLD_AND_BUY_TEMPLATE.format(_p_created_date=datetime.datetime.now(),
                                                _p_rebalance_interval=duration,
                                                _p_max_holding_num=hold_cnt,
                                                _p_single_profit_taken=gain_limit1,
                                                _p_single_profit_max_drawdown=gain_limit2,
                                                _p_single_stop_loss=loss_limit)
        elif policy_type == PolicyClass.orders:
            code = ORDERS_TEMPLATE.format(_p_created_date=datetime.datetime.now())
        elif policy_type == PolicyClass.rebalance:
            code = REBALANCE_TEMPLATE.format(_p_created_date=datetime.datetime.now())
        elif policy_type == PolicyClass.hold_spec_days:
            code = HOLD_SPEC_DAYS_TEMPLATE.format(_p_created_date=datetime.datetime.now(),
                                                  _p_rebalance_interval=duration)
        else:
            raise ValueError("Unknown policy {}".format(policy_type))
        return duration, code, bt_config.copy()
